from pandas import DataFrame
import pandas as pd

from openpyxl.utils import get_column_letter
from pandas import ExcelWriter
import numpy as np

from .sqlite_db import create_conn


def get_cols_values_dict(df: DataFrame, col_name_list):
    """
    获取DataFrame中的列的值得字典
    {col_name:[list], ---}
    """
    need_data_df = df[col_name_list]
    return need_data_df.to_dict(orient="list")


# region 数据库操作
def get_query_df(sql_str, count=500) -> DataFrame:
    conn = create_conn()
    """获取查询语句中的数据"""
    df_data = pd.read_sql(sql_str, conn).head(count)
    conn.close()
    return df_data


def get_df_from_db(table_name, condition, count=500):
    sql_str = f"select * from {table_name}"
    if condition:
        sql_str = sql_str + f" where {condition}"
    return get_query_df(sql_str, count)


# endregion

# region Excel 操作
def save_df_to_excel(df: DataFrame, excel_path, sheet_name="sheet_1"):
    with pd.ExcelWriter(excel_path) as writer:
        to_excel_auto_column_weight(df, writer, sheet_name)
        return True


def to_excel_auto_column_weight(df: pd.DataFrame, writer: ExcelWriter, sheet_name):
    """DataFrame保存为excel并自动设置列宽"""
    df.to_excel(writer, sheet_name=sheet_name, index=False)
    #  计算表头的字符宽度
    column_widths = (
        df.columns.to_series().apply(lambda x: len(x.encode('gbk'))).values
    )
    #  计算每列的最大字符宽度
    max_widths = (
        df.astype(str).applymap(lambda x: len(x.encode('gbk'))).agg(max).values
    )
    # 计算整体最大宽度
    widths = np.max([column_widths, max_widths], axis=0)
    # 设置列宽
    worksheet = writer.sheets[sheet_name]
    for i, width in enumerate(widths, 1):
        # openpyxl引擎设置字符宽度时会缩水0.5左右个字符，所以干脆+2使左右都空出一个字宽。
        worksheet.column_dimensions[get_column_letter(i)].width = width + 2


# endregion

# region pandas 功能

# endregion
